Use xp.stack in messages so JAX can trace fixed arrays - #1460
Merged
Conversation
Jammy2211
marked this pull request as ready for review
August 9, 2026 16:21
This was referenced Aug 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
xp.array([...])message constructors withxp.stack([...])GammaMessage.to_canonical_formthroughxp.loginstead of NumPyjax.jitcoverage for every changed constructor, with NumPy shape/value parityWhy
The Gamma canonical-form path currently sends a traced value through
np.log, raisingTracerArrayConversionError. Using the selected array namespace throughout keeps message sufficient-statistic construction traceable and makes the fixed-array construction consistent across NumPy and JAX.The same regression test was run against untouched
main: the scalar and batched Gamma canonical-form cases fail there and pass on this branch. The other array constructors already pass with current supported JAX, so theirxp.stackconversion is cross-backend hardening with unchanged shapes and values.Verification
test_autofit/messages: 32 passedtest_autofit: 1689 passed, 4 skippedmain, bothGammaMessage.to_canonical_form; 20 passed on this branchThe private
z_projects/concrintegration reproducer was not available in this environment.Prior/message audit
The adjacent audit found four separate JIT backend leaks in Gamma/Beta
log_partitionand compound-priorlog/log10. They are independently reproduced and tracked in #1459 rather than widening this PR.Closes #1458